import sys
from pathlib import Path

# Add project root to sys.path to allow running tests from this directory
# The project root is 6 levels up from the parent directory of this file.
project_root = str(Path(__file__).resolve().parents[6])
if project_root not in sys.path:
    sys.path.insert(0, project_root)

#
# Copyright © 2024 Agora
# This file is part of TEN Framework, an open source project.
# Licensed under the Apache License, Version 2.0, with certain conditions.
# Refer to the "LICENSE" file in the root directory for more information.
#
import json
from typing import Any
from unittest.mock import patch, AsyncMock
import tempfile
import os
import asyncio
import filecmp
import shutil
import threading

from ten_runtime import (
    ExtensionTester,
    TenEnvTester,
    Cmd,
    CmdResult,
    StatusCode,
    Data,
)
from ten_ai_base.struct import TTSTextInput, TTSFlush
from ten_ai_base.message import ModuleVendorException, ModuleErrorVendorInfo


# ================ test dump ================
class ExtensionTesterDump(ExtensionTester):
    def __init__(self):
        super().__init__()
        # Use a fixed path as requested by the user.
        self.dump_dir = "./dump/"
        # Use a unique name for the file generated by the test to avoid collision
        # with the file generated by the extension.
        self.test_dump_file_path = os.path.join(
            self.dump_dir, "test_manual_dump.pcm"
        )
        self.audio_end_received = False
        self.received_audio_chunks = []

    def on_start(self, ten_env_tester: TenEnvTester) -> None:
        """Called when test starts, sends a TTS request."""
        ten_env_tester.log_info("Dump test started, sending TTS request.")

        tts_input = TTSTextInput(
            request_id="tts_request_1",
            text="hello word, hello agora",
            text_input_end=True,
        )
        data = Data.create("tts_text_input")
        data.set_property_from_json(None, tts_input.model_dump_json())
        ten_env_tester.send_data(data)
        ten_env_tester.on_start_done()

    def on_data(self, ten_env: TenEnvTester, data) -> None:
        name = data.get_name()
        if name == "tts_audio_end":
            ten_env.log_info("Received tts_audio_end, stopping test.")
            self.audio_end_received = True
            ten_env.stop_test()

    def on_audio_frame(self, ten_env: TenEnvTester, audio_frame):
        """Receives audio frames and collects their data using the lock/unlock pattern."""
        # The 'audio_frame' object is a wrapper around a memory buffer.
        # We must lock the buffer to safely access the data, copy it,
        # and finally unlock the buffer so the runtime can reuse it.
        buf = audio_frame.lock_buf()
        try:
            # We must copy the data from the buffer, as the underlying memory
            # may be freed or reused after we unlock it.
            copied_data = bytes(buf)
            self.received_audio_chunks.append(copied_data)
        finally:
            # Always ensure the buffer is unlocked, even if an error occurs.
            audio_frame.unlock_buf(buf)

    def write_test_dump_file(self):
        """Writes the collected audio chunks to a file."""
        with open(self.test_dump_file_path, "wb") as f:
            for chunk in self.received_audio_chunks:
                f.write(chunk)

    def find_tts_dump_file(self) -> str | None:
        """Find the dump file created by the TTS extension in the fixed dump directory."""
        if not os.path.exists(self.dump_dir):
            return None
        for filename in os.listdir(self.dump_dir):
            if filename.endswith(".pcm") and filename != os.path.basename(
                self.test_dump_file_path
            ):
                return os.path.join(self.dump_dir, filename)
        return None


@patch("google_tts_python.extension.GoogleTTS")
def test_dump_functionality(MockGoogleTTS):
    """Test that the extension can dump audio to a file."""
    # Create dump directory if it doesn't exist
    dump_dir = "./dump/"
    os.makedirs(dump_dir, exist_ok=True)

    # Mock the GoogleTTS class
    mock_client_instance = AsyncMock()

    # Mock the get method to return audio data as async generator
    async def mock_get(text, request_id):
        # Return audio data in the expected format for streaming
        yield b"fake_audio_data_1", 1, 123  # EVENT_TTS_RESPONSE with ttfb_ms
        yield b"fake_audio_data_2", 1, None  # EVENT_TTS_RESPONSE without ttfb_ms
        yield None, 2, None  # EVENT_TTS_REQUEST_END

    # Set the mock methods
    mock_client_instance.get = mock_get
    mock_client_instance.get_streaming = mock_get  # Streaming method
    mock_client_instance.clean = AsyncMock()
    mock_client_instance.reset = AsyncMock()
    mock_client_instance._should_stop_streaming = False

    # Set up required attributes
    mock_client_instance.client = AsyncMock()
    mock_client_instance.config = AsyncMock()
    mock_client_instance.ten_env = AsyncMock()
    mock_client_instance._is_cancelled = False
    mock_client_instance.credentials = None
    mock_client_instance.send_text_in_connection = False
    mock_client_instance.cur_request_id = ""
    mock_client_instance.streaming_enabled = False
    mock_client_instance.streaming_config = None

    # Mock the constructor to return our mock instance
    MockGoogleTTS.return_value = mock_client_instance

    # Mock the _initialize_client method to avoid actual initialization
    mock_client_instance._initialize_client = AsyncMock()

    # Create and run the tester
    tester = ExtensionTesterDump()

    # Set up dump configuration
    dump_config = {
        "dump": True,
        "dump_path": dump_dir,
        "params": {
            "AudioConfig": {"sample_rate_hertz": 16000},
            "VoiceSelectionParams": {
                "language_code": "en-US",
                "ssml_gender": "NEUTRAL",
            },
            "credentials": "fake_credentials_for_mock_testing",
        },
    }

    tester.set_test_mode_single("google_tts_python", json.dumps(dump_config))
    tester.run()

    # Verify that audio end was received
    assert tester.audio_end_received, "Audio end event was not received"

    # Write test dump file for comparison
    tester.write_test_dump_file()


# ================ test text input end ================
class ExtensionTesterTextInputEnd(ExtensionTester):
    def __init__(self):
        super().__init__()
        self.ten_env: TenEnvTester | None = None
        self.first_request_completed = False
        self.second_request_sent = False
        self.audio_end_count = 0

    def on_start(self, ten_env_tester: TenEnvTester) -> None:
        """Called when test starts, sends first TTS request."""
        self.ten_env = ten_env_tester
        ten_env_tester.log_info(
            "Text input end test started, sending first TTS request."
        )

        tts_input = TTSTextInput(
            request_id="tts_request_1",
            text="hello word, hello agora",
            text_input_end=True,
        )
        data = Data.create("tts_text_input")
        data.set_property_from_json(None, tts_input.model_dump_json())
        ten_env_tester.send_data(data)
        ten_env_tester.on_start_done()

    def on_data(self, ten_env: TenEnvTester, data) -> None:
        name = data.get_name()
        if name == "tts_audio_end":
            self.audio_end_count += 1
            ten_env.log_info(
                f"Received tts_audio_end, count: {self.audio_end_count}"
            )

            if self.audio_end_count == 1 and not self.second_request_sent:
                self.first_request_completed = True
                self.second_request_sent = True
                ten_env.log_info(
                    "First request completed, sending second request."
                )

                # Send second request
                tts_input = TTSTextInput(
                    request_id="tts_request_2",
                    text="second request",
                    text_input_end=True,
                )
                data = Data.create("tts_text_input")
                data.set_property_from_json(None, tts_input.model_dump_json())
                ten_env.send_data(data)
            elif self.audio_end_count == 2:
                ten_env.log_info("Second request completed, stopping test.")
                ten_env.stop_test()


@patch("google_tts_python.extension.GoogleTTS")
def test_text_input_end(MockGoogleTTS):
    """Test that the extension handles text_input_end correctly."""
    # Mock the GoogleTTS class
    mock_client_instance = AsyncMock()

    # Mock the get method to return audio data for both requests
    async def mock_get(text, request_id):
        # Return audio data in the expected format for streaming
        yield b"fake_audio_data", 1, 123  # EVENT_TTS_RESPONSE with ttfb_ms
        yield None, 2, None  # EVENT_TTS_REQUEST_END

    mock_client_instance.get = mock_get
    mock_client_instance.get_streaming = mock_get  # Streaming method
    mock_client_instance.clean = AsyncMock()
    mock_client_instance.reset = AsyncMock()
    mock_client_instance._should_stop_streaming = False

    # Set up required attributes
    mock_client_instance.client = AsyncMock()
    mock_client_instance.config = AsyncMock()
    mock_client_instance.ten_env = AsyncMock()
    mock_client_instance._is_cancelled = False
    mock_client_instance.credentials = None
    mock_client_instance.send_text_in_connection = False
    mock_client_instance.cur_request_id = ""
    mock_client_instance.streaming_enabled = False
    mock_client_instance.streaming_config = None

    # Mock the constructor to return our mock instance
    MockGoogleTTS.return_value = mock_client_instance

    # Mock the _initialize_client method to avoid actual initialization
    mock_client_instance._initialize_client = AsyncMock()

    # Create and run the tester
    tester = ExtensionTesterTextInputEnd()

    # Set up configuration
    config = {
        "params": {
            "AudioConfig": {"sample_rate_hertz": 16000},
            "VoiceSelectionParams": {
                "language_code": "en-US",
                "ssml_gender": "NEUTRAL",
            },
            "credentials": "fake_credentials_for_mock_testing",
        },
    }

    tester.set_test_mode_single("google_tts_python", json.dumps(config))
    tester.run()

    # Verify that both requests completed
    assert (
        tester.audio_end_count == 2
    ), f"Expected 2 audio end events, got {tester.audio_end_count}"


# ================ test flush ================
class ExtensionTesterFlush(ExtensionTester):
    def __init__(self):
        super().__init__()
        self.audio_start_received = False
        self.flush_sent = False
        self.audio_end_received = False

    def on_start(self, ten_env_tester: TenEnvTester) -> None:
        """Called when test starts, sends a TTS request."""
        ten_env_tester.log_info("Flush test started, sending TTS request.")

        tts_input = TTSTextInput(
            request_id="tts_request_1",
            text="hello word, hello agora",
            text_input_end=True,
        )
        data = Data.create("tts_text_input")
        data.set_property_from_json(None, tts_input.model_dump_json())
        ten_env_tester.send_data(data)
        ten_env_tester.on_start_done()

    def on_data(self, ten_env: TenEnvTester, data) -> None:
        name = data.get_name()
        if name == "tts_audio_start":
            ten_env.log_info("Received tts_audio_start, sending flush.")
            self.audio_start_received = True

            # Send flush request
            flush_data = Data.create("tts_flush")
            flush_data.set_property_string("flush_id", "tts_request_1")
            ten_env.send_data(flush_data)
            self.flush_sent = True

        elif name == "tts_audio_end":
            ten_env.log_info("Received tts_audio_end, stopping test.")
            self.audio_end_received = True
            ten_env.stop_test()


@patch("google_tts_python.extension.GoogleTTS")
def test_flush_functionality(MockGoogleTTS):
    """Test that the extension handles flush correctly."""
    # Mock the GoogleTTS class
    mock_client_instance = AsyncMock()

    # Mock the get method to return audio data with flush support
    async def mock_get(text, request_id):
        # Return audio data in the expected format for streaming
        yield b"fake_audio_data", 1, 123  # EVENT_TTS_RESPONSE with ttfb_ms
        # After flush, should not continue - this is handled by the _should_stop_streaming flag
        # The mock will stop yielding when clean() is called

    mock_client_instance.get = mock_get
    mock_client_instance.get_streaming = mock_get  # Streaming method
    mock_client_instance.clean = AsyncMock()
    mock_client_instance.reset = AsyncMock()
    mock_client_instance._should_stop_streaming = False

    # Set up required attributes
    mock_client_instance.client = AsyncMock()
    mock_client_instance.config = AsyncMock()
    mock_client_instance.ten_env = AsyncMock()
    mock_client_instance._is_cancelled = False
    mock_client_instance.credentials = None
    mock_client_instance.send_text_in_connection = False
    mock_client_instance.cur_request_id = ""
    mock_client_instance.streaming_enabled = False
    mock_client_instance.streaming_config = None

    # Mock the constructor to return our mock instance
    MockGoogleTTS.return_value = mock_client_instance

    # Mock the _initialize_client method to avoid actual initialization
    mock_client_instance._initialize_client = AsyncMock()

    # Create and run the tester
    tester = ExtensionTesterFlush()

    # Set up configuration
    config = {
        "params": {
            "AudioConfig": {"sample_rate_hertz": 16000},
            "VoiceSelectionParams": {
                "language_code": "en-US",
                "ssml_gender": "NEUTRAL",
            },
            "credentials": "fake_credentials_for_mock_testing",
        },
    }

    tester.set_test_mode_single("google_tts_python", json.dumps(config))
    tester.run()

    # Verify that flush was handled correctly
    assert tester.audio_start_received, "Audio start event was not received"
    assert tester.flush_sent, "Flush was not sent"
    assert tester.audio_end_received, "Audio end event was not received"


# ================ test error handling ================
class ExtensionTesterError(ExtensionTester):
    def __init__(self):
        super().__init__()
        self.error_received = False

    def on_start(self, ten_env_tester: TenEnvTester) -> None:
        """Called when test starts, sends a TTS request."""
        ten_env_tester.log_info("Error test started, sending TTS request.")

        tts_input = TTSTextInput(
            request_id="tts_request_1",
            text="hello word, hello agora",
            text_input_end=True,
        )
        data = Data.create("tts_text_input")
        data.set_property_from_json(None, tts_input.model_dump_json())
        ten_env_tester.send_data(data)
        ten_env_tester.on_start_done()

    def on_data(self, ten_env: TenEnvTester, data) -> None:
        name = data.get_name()
        if name == "error":
            ten_env.log_info("Received error, stopping test.")
            self.error_received = True
            ten_env.stop_test()


@patch("google_tts_python.extension.GoogleTTS")
def test_error_handling(MockGoogleTTS):
    """Test that the extension handles errors correctly."""
    # Mock the GoogleTTS class to raise an exception
    mock_client_instance = AsyncMock()

    # Mock the get method to raise an exception
    async def mock_get(text, request_id):
        raise Exception("Test error")

    mock_client_instance.get = mock_get
    mock_client_instance.get_streaming = mock_get  # Streaming method
    mock_client_instance.clean = AsyncMock()
    mock_client_instance.reset = AsyncMock()
    mock_client_instance._should_stop_streaming = False

    # Set up required attributes
    mock_client_instance.client = AsyncMock()
    mock_client_instance.config = AsyncMock()
    mock_client_instance.ten_env = AsyncMock()
    mock_client_instance._is_cancelled = False
    mock_client_instance.credentials = None
    mock_client_instance.send_text_in_connection = False
    mock_client_instance.cur_request_id = ""
    mock_client_instance.streaming_enabled = False
    mock_client_instance.streaming_config = None

    # Mock the constructor to return our mock instance
    MockGoogleTTS.return_value = mock_client_instance

    # Mock the _initialize_client method to avoid actual initialization
    mock_client_instance._initialize_client = AsyncMock()

    # Create and run the tester
    tester = ExtensionTesterError()

    # Set up configuration
    config = {
        "params": {
            "AudioConfig": {"sample_rate_hertz": 16000},
            "VoiceSelectionParams": {
                "language_code": "en-US",
                "ssml_gender": "NEUTRAL",
            },
            "credentials": "fake_credentials_for_mock_testing",
        },
    }

    tester.set_test_mode_single("google_tts_python", json.dumps(config))
    tester.run()

    # Verify that error was handled correctly
    assert tester.error_received, "Error event was not received"


# ================ test basic functionality ================
class ExtensionTesterBasic(ExtensionTester):
    def __init__(self):
        super().__init__()
        self.audio_end_received = False

    def on_start(self, ten_env_tester: TenEnvTester) -> None:
        """Called when test starts, sends a TTS request."""
        ten_env_tester.log_info("Basic test started, sending TTS request.")

        tts_input = TTSTextInput(
            request_id="tts_request_1",
            text="hello word, hello agora",
            text_input_end=True,
        )
        data = Data.create("tts_text_input")
        data.set_property_from_json(None, tts_input.model_dump_json())
        ten_env_tester.send_data(data)
        ten_env_tester.on_start_done()

    def on_data(self, ten_env: TenEnvTester, data) -> None:
        name = data.get_name()
        if name == "tts_audio_end":
            ten_env.log_info("Received tts_audio_end, stopping test.")
            self.audio_end_received = True
            ten_env.stop_test()


@patch("google_tts_python.extension.GoogleTTS")
def test_basic_functionality(MockGoogleTTS):
    """Test basic TTS functionality."""
    # Mock the GoogleTTS class
    mock_client_instance = AsyncMock()

    # Mock the get method to return audio data
    async def mock_get(text, request_id):
        # Return audio data in the expected format for streaming
        yield b"fake_audio_data", 1, 123  # EVENT_TTS_RESPONSE with ttfb_ms
        yield None, 2, None  # EVENT_TTS_REQUEST_END

    mock_client_instance.get = mock_get
    mock_client_instance.get_streaming = mock_get  # Streaming method
    mock_client_instance.clean = AsyncMock()
    mock_client_instance.reset = AsyncMock()
    mock_client_instance._should_stop_streaming = False

    # Set up required attributes
    mock_client_instance.client = AsyncMock()
    mock_client_instance.config = AsyncMock()
    mock_client_instance.ten_env = AsyncMock()
    mock_client_instance._is_cancelled = False
    mock_client_instance.credentials = None
    mock_client_instance.send_text_in_connection = False
    mock_client_instance.cur_request_id = ""
    mock_client_instance.streaming_enabled = False
    mock_client_instance.streaming_config = None

    # Mock the constructor to return our mock instance
    MockGoogleTTS.return_value = mock_client_instance

    # Mock the _initialize_client method to avoid actual initialization
    mock_client_instance._initialize_client = AsyncMock()

    # Create and run the tester
    tester = ExtensionTesterBasic()

    # Set up configuration
    config = {
        "params": {
            "AudioConfig": {"sample_rate_hertz": 16000},
            "VoiceSelectionParams": {
                "language_code": "en-US",
                "ssml_gender": "NEUTRAL",
            },
            "credentials": "fake_credentials_for_mock_testing",
        },
    }

    tester.set_test_mode_single("google_tts_python", json.dumps(config))
    tester.run()

    # Verify that audio end was received
    assert tester.audio_end_received, "Audio end event was not received"
